"""
视频显示组件模块

该模块实现了视频显示和目标选择功能。
主要功能：
1. 视频帧显示
2. 鼠标交互目标选择
3. 图像缩放和适配
4. 选择框绘制

作者：目标跟踪系统
日期：2024
"""

import cv2
import numpy as np
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from typing import Optional, Tuple


class VideoWidget(QLabel):
    """
    视频显示组件类
    
    继承自QLabel，实现视频帧显示和目标选择功能。
    支持鼠标拖拽选择目标区域。
    """
    
    # 自定义信号
    target_selected = pyqtSignal(tuple)  # 目标选择信号
    
    def __init__(self, parent=None):
        super().__init__(parent)
        
        # 设置基本属性
        self.setMinimumSize(640, 480)
        self.setStyleSheet("border: 1px solid gray")
        self.setAlignment(Qt.AlignCenter)
        self.setText("请加载视频文件")
        
        # 鼠标交互状态
        self.selecting = False
        self.start_point = None
        self.end_point = None
        self.selection_rect = None
        
        # 图像相关
        self.original_frame = None
        self.display_frame = None
        self.scale_factor = 1.0
        self.offset_x = 0
        self.offset_y = 0
        
    def update_frame(self, frame: np.ndarray):
        """
        更新显示帧
        
        参数：
            frame: 输入帧（BGR格式）
        """
        if frame is None:
            return
            
        self.original_frame = frame.copy()
        
        # 转换为RGB格式
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # 计算缩放参数
        self._calculate_scale_params(rgb_frame)
        
        # 缩放图像
        scaled_frame = self._scale_frame(rgb_frame)
        
        # 转换为QImage
        height, width, channel = scaled_frame.shape
        bytes_per_line = 3 * width
        q_image = QImage(scaled_frame.data, width, height, bytes_per_line, QImage.Format_RGB888)
        
        # 转换为QPixmap并显示
        pixmap = QPixmap.fromImage(q_image)
        self.setPixmap(pixmap)
        
        self.display_frame = scaled_frame
        
    def _calculate_scale_params(self, frame: np.ndarray):
        """
        计算缩放参数
        
        参数：
            frame: 输入帧
        """
        widget_width = self.width()
        widget_height = self.height()
        frame_height, frame_width = frame.shape[:2]
        
        # 计算缩放比例（保持宽高比）
        scale_x = widget_width / frame_width
        scale_y = widget_height / frame_height
        self.scale_factor = min(scale_x, scale_y)
        
        # 计算偏移量（居中显示）
        scaled_width = int(frame_width * self.scale_factor)
        scaled_height = int(frame_height * self.scale_factor)
        
        self.offset_x = (widget_width - scaled_width) // 2
        self.offset_y = (widget_height - scaled_height) // 2
        
    def _scale_frame(self, frame: np.ndarray) -> np.ndarray:
        """
        缩放帧到合适尺寸
        
        参数：
            frame: 输入帧
            
        返回：
            缩放后的帧
        """
        frame_height, frame_width = frame.shape[:2]
        new_width = int(frame_width * self.scale_factor)
        new_height = int(frame_height * self.scale_factor)
        
        if new_width > 0 and new_height > 0:
            return cv2.resize(frame, (new_width, new_height))
        else:
            return frame
            
    def mousePressEvent(self, event):
        """
        鼠标按下事件
        """
        if event.button() == Qt.LeftButton and self.original_frame is not None:
            # 转换坐标到原始图像坐标系
            widget_x = event.x()
            widget_y = event.y()
            
            # 检查是否在图像区域内
            if self._is_in_image_area(widget_x, widget_y):
                self.selecting = True
                self.start_point = (widget_x, widget_y)
                self.end_point = (widget_x, widget_y)
                
    def mouseMoveEvent(self, event):
        """
        鼠标移动事件
        """
        if self.selecting:
            widget_x = event.x()
            widget_y = event.y()
            
            # 限制在图像区域内
            widget_x = max(self.offset_x, min(widget_x, self.width() - self.offset_x))
            widget_y = max(self.offset_y, min(widget_y, self.height() - self.offset_y))
            
            self.end_point = (widget_x, widget_y)
            self.update()  # 触发重绘
            
    def mouseReleaseEvent(self, event):
        """
        鼠标释放事件
        """
        if event.button() == Qt.LeftButton and self.selecting:
            self.selecting = False
            
            if self.start_point and self.end_point:
                # 计算选择区域
                bbox = self._calculate_selection_bbox()
                
                if bbox and bbox[2] > 10 and bbox[3] > 10:  # 最小尺寸检查
                    self.target_selected.emit(bbox)
                    
            self.start_point = None
            self.end_point = None
            self.update()
            
    def paintEvent(self, event):
        """
        绘制事件
        """
        super().paintEvent(event)
        
        # 绘制选择框
        if self.selecting and self.start_point and self.end_point:
            painter = QPainter(self)
            painter.setPen(QPen(Qt.red, 2, Qt.SolidLine))
            
            x1, y1 = self.start_point
            x2, y2 = self.end_point
            
            rect = QRect(min(x1, x2), min(y1, y2), abs(x2 - x1), abs(y2 - y1))
            painter.drawRect(rect)
            
    def _is_in_image_area(self, x: int, y: int) -> bool:
        """
        检查坐标是否在图像显示区域内
        
        参数：
            x: X坐标
            y: Y坐标
            
        返回：
            是否在图像区域内
        """
        if self.display_frame is None:
            return False
            
        height, width = self.display_frame.shape[:2]
        
        return (self.offset_x <= x <= self.offset_x + width and 
                self.offset_y <= y <= self.offset_y + height)
                
    def _calculate_selection_bbox(self) -> Optional[Tuple[int, int, int, int]]:
        """
        计算选择区域的边界框（原始图像坐标系）
        
        返回：
            边界框 (x, y, width, height)，如果无效则返回None
        """
        if not self.start_point or not self.end_point or self.original_frame is None:
            return None
            
        # 组件坐标
        x1, y1 = self.start_point
        x2, y2 = self.end_point
        
        # 确保左上角为起点
        widget_x = min(x1, x2)
        widget_y = min(y1, y2)
        widget_w = abs(x2 - x1)
        widget_h = abs(y2 - y1)
        
        # 转换为图像坐标
        image_x = int((widget_x - self.offset_x) / self.scale_factor)
        image_y = int((widget_y - self.offset_y) / self.scale_factor)
        image_w = int(widget_w / self.scale_factor)
        image_h = int(widget_h / self.scale_factor)
        
        # 边界检查
        frame_height, frame_width = self.original_frame.shape[:2]
        
        image_x = max(0, min(image_x, frame_width - 1))
        image_y = max(0, min(image_y, frame_height - 1))
        image_w = min(image_w, frame_width - image_x)
        image_h = min(image_h, frame_height - image_y)
        
        if image_w > 0 and image_h > 0:
            return (image_x, image_y, image_w, image_h)
        else:
            return None
            
    def wheelEvent(self, event):
        """
        鼠标滚轮事件（预留用于缩放功能）
        """
        # 可以在此实现图像缩放功能
        pass
        
    def reset(self):
        """
        重置显示状态
        """
        self.original_frame = None
        self.display_frame = None
        self.selecting = False
        self.start_point = None
        self.end_point = None
        self.selection_rect = None
        
        self.clear()
        self.setText("请加载视频文件")
        
    def get_display_size(self) -> Tuple[int, int]:
        """
        获取当前显示尺寸
        
        返回：
            (宽度, 高度)
        """
        if self.display_frame is not None:
            height, width = self.display_frame.shape[:2]
            return (width, height)
        else:
            return (0, 0)
            
    def get_original_size(self) -> Tuple[int, int]:
        """
        获取原始图像尺寸
        
        返回：
            (宽度, 高度)
        """
        if self.original_frame is not None:
            height, width = self.original_frame.shape[:2]
            return (width, height)
        else:
            return (0, 0) 